package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"golang.org/x/crypto/ssh"
	"io"
	"log"
	"net"
	"net/url"
	"os"
	"strings"
	"time"
)

type Config struct {
	User      string `json:"user"`       //ssh username
	Pass      string `json:"pass"`       //ssh password
	HostAddr  string `json:"host_addr"`  //ssh host:port
	LocalAddr string `json:"local_addr"` //localhost:port
}

var cfg struct {
	Debug     bool   `json:"debug"`
	User      string `json:"user"`       //ssh username
	Pass      string `json:"pass"`       //ssh password
	HostAddr  string `json:"host_addr"`  //ssh host:port
	LocalAddr string `json:"local_addr"` //localhost:port
}

func main() {
	err := loadConfig()
	if err != nil {
		log.Fatalf("load config err: %v", err)
	}
	config := &ssh.ClientConfig{
		User: cfg.User,
		Auth: []ssh.AuthMethod{
			ssh.Password(cfg.Pass),
		},
		Timeout: 30 * time.Second,
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}
	sshConn, err := ssh.Dial("tcp", cfg.HostAddr, config)
	if err != nil {
		log.Fatalf("ssh.Dial failed: %#v", err)
	}
	defer sshConn.Close()

	listener, err := net.Listen("tcp", cfg.LocalAddr)
	if cfg.Debug {
		log.Printf("listen:%s\n", cfg.LocalAddr)
	}

	if err != nil {
		log.Fatalf("net.Listen failed: %#v", err)
	}
	if cfg.Debug {
		log.Printf("net.Listen success: %#v\n", err)
	}
	for {
		// Setup localConn (type net.Conn)
		localConn, err := listener.Accept()
		if err != nil {
			log.Fatalf("listen.Accept failed: %v", err)
		}
		if cfg.Debug {
			log.Printf("Access Accepted\n")
		}
		go forward(&localConn, sshConn)
	}
}

func loadConfig() (err error) {
	path := "./config.json"
	_, err = os.Stat(path)
	if err != nil {
		err = errors.New("config.json not exists")
		return
	}
	f, err := os.Open(path)
	err = json.NewDecoder(f).Decode(&cfg)
	return
}

func forward(localConn *net.Conn, sshConn *ssh.Client) {
	//获取本地客户端请求
	address, method, buf, err := GetRequestAddr(localConn)
	if err != nil {
		fmt.Printf("[line:36]read local request err: %#v\n", err)
		return
	}
	//获得了请求的 host 和 port，代理服务器向目标服务端发起 tcp 连接
	var server net.Conn
	server, err = sshConn.Dial("tcp", address)
	if err != nil {
		fmt.Printf("[line:43] ssh client request err: %#v\n", err)
		return
	}
	defer server.Close()
	//如果使用 https 协议，需先向客户端表示连接建立完毕
	if method == "CONNECT" {
		fmt.Fprint(*localConn, "HTTP/1.1 200 Connection established\r\n\r\n")
	} else { //如果使用 http 协议，需将从客户端得到的 http 请求转发给服务端
		server.Write(buf)
	}
	//将客户端的请求转发至服务端，将服务端的响应转发给客户端。io.Copy 为阻塞函数，文件描述符不关闭就不停止
	ch := make(chan error, 1)
	go CopyTo(ch, server, *localConn)
	go CopyTo(ch, *localConn, server)
	err = <-ch
}

func CopyTo(e chan<- error, dst, src net.Conn) {
	defer func() {
		recover()
	}()
	_, err := io.Copy(dst, src)
	e <- err
}

func GetRequestAddr(client *net.Conn) (address, method string, proto []byte, err error) {
	// 用来存放客户端数据的缓冲区
	var b [1024]byte
	//从客户端获取数据
	var n int
	n, err = (*client).Read(b[:])
	if err != nil {
		log.Printf("Request msg read fail[120]: %#v\n", err)
		return
	}
	if cfg.Debug {
		log.Printf("Request Mesage:\n%s\n", b[:n])
	}
	var URL string
	// 从客户端数据读入 method，url
	buf := b[:bytes.IndexByte(b[:], '\n')]
	fmt.Sscanf(string(buf), "%s%s", &method, &URL)
	if cfg.Debug {
		log.Printf("client request--------->\n%s\n", string(buf))
	}
	var hostPortURL *url.URL
	hostPortURL, err = url.Parse(URL)
	if err != nil {
		log.Printf("url parse fail[132]:%#v\n", err)
		return
	}
	// 如果方法是 CONNECT，则为 https 协议
	if method == "CONNECT" {
		address = hostPortURL.Scheme + ":" + hostPortURL.Opaque
	} else { //否则为 http 协议
		address = hostPortURL.Host
		// 如果 host 不带端口，则默认为 80
		if strings.Index(hostPortURL.Host, ":") == -1 { //host 不带端口， 默认 80
			address = hostPortURL.Host + ":80"
		}
		proto = b[:n]
	}
	return
}
